from xml.dom.minidom import Document


def add_child_branch(name, parent_branch, doc, attribute_name_list,
                     attribute_value_list):
    child_branch = doc.createElement(name)
    max_length = max(len(attribute_name_list), len(attribute_value_list))
    for i in range(max_length):
        if i >= len(attribute_value_list):
            attribute_value_list.append(0)
        if i >= len(attribute_name_list):
            attribute_name_list.append("default")

    for i in range(max_length):
        child_branch.setAttribute(attribute_name_list[i],
                                  attribute_value_list[i])

    parent_branch.appendChild(child_branch)
    return child_branch


class Link:

    def __init__(self, root, doc, link_name, link_width, link_length,
                 joint_name, joint_type, axis):  # 初始化方法
        self.root = root
        self.doc = doc
        self.link_width = link_width
        self.link_length = link_length
        self.joint_name = joint_name
        self.joint_type = joint_type
        self.axis = axis
        self.link = add_child_branch('link', root, doc, ['name'], [link_name])

    def add_part(self, part_name):
        size = str(self.link_width) + " " + str(self.link_width) + " " + str(
            self.link_length)
        ori = str(0) + " " + str(0) + " " + str(self.link_length / 2)

        part = add_child_branch(part_name, self.link, self.doc, [], [])
        geometry = add_child_branch('geometry', part, self.doc, [], [])

        add_child_branch('box', geometry, self.doc, ['size'], [size])
        add_child_branch('origin', part, self.doc, ['xyz', 'rpy'],
                         [ori, '0 0 0'])

    def add_joint(self, parent_link, child_link):
        ori = str(0) + " " + str(0) + " " + str(self.link_length)

        if self.axis == 'x':
            axis_round = '1 0 0'
        elif self.axis == 'y':
            axis_round = '0 1 0'
        elif self.axis == 'z':
            axis_round = '0 0 1'
        else:
            return False

        joint = add_child_branch('joint', self.root, self.doc,
                                 ['name', 'type'],
                                 [self.joint_name, self.joint_type])

        add_child_branch('axis', joint, self.doc, ['xyz'], [axis_round])
        add_child_branch('parent', joint, self.doc, ['link'], [parent_link])
        add_child_branch('child', joint, self.doc, ['link'], [child_link])
        add_child_branch('origin', joint, self.doc, ['xyz', 'rpy'],
                         [ori, '0 0 0'])

        if self.joint_type == 'revolute':
            limit = self.doc.createElement('limit')
            limit.setAttribute('effort', '1000')
            limit.setAttribute('lower', '-3.14')
            limit.setAttribute('upper', '3.14159')
            limit.setAttribute('velocity', '0')
            joint.appendChild(limit)

        return True


class Robot:

    def __init__(self, name):  # 初始化方法
        self.doc = Document()
        self.root = add_child_branch('robot', self.doc, self.doc, ['name'],
                                     [name])

    def add_link(self,
                 name,
                 child_link,
                 width,
                 length,
                 joint_name,
                 joint_type,
                 axis,
                 final=0):
        link = Link(self.root, self.doc, name, width, length, joint_name,
                    joint_type, axis)
        link.add_part('visual')
        link.add_part('collision')
        if final == 0:
            link.add_joint(name, child_link)

    def output(self, path):
        with open(path, "w+") as f:
            # 文件转码decode是为了保证输出的是UTF-8格式
            f.write(self.doc.toprettyxml(encoding="UTF-8").decode("UTF-8"))
            # doc.writexml(f, indent='	', add_indent='	', new_line='', encoding='utf-8')
            f.close()
